'''
Description: 
Author: 爱酱丶
Date: 2022-08-09 17:30:04
LastEditTime: 2022-10-10 16:56:39
LastEditors: 爱酱丶
'''
from osgeo import gdal
from osgeo import ogr
from osgeo import osr
from osgeo import gdalconst
import os

# transformation=osr.CreateCoordinateTransformation(spatialrefMode,spatialrefOld)
# geomTransformer=ogr.GeomTransformer(transformation)
# print(geometry.ExportToWkt())
# geometry=geomTransformer.Transform(geometry)
# print(geometry.ExportToWkt())

def vector_erase(*args):
    old_path, mode_path, fieldName, fieldValue, new_path = args
    gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")  
    datasourceOld=ogr.Open(old_path)
    driver = datasourceOld.GetDriver()
    layerOld=datasourceOld.GetLayer(0)
    layerName=layerOld.GetName()
    spatialrefOld = layerOld.GetSpatialRef()
    geomType=layerOld.GetGeomType()
    schema=layerOld.schema
    print(layerOld.GetFeatureCount())

    for i in schema:print(i)

    datasourceMode=ogr.Open(mode_path)
    layerMode=datasourceMode.GetLayer()
    layerMode.SetAttributeFilter(f"{fieldName}='{fieldValue}'")

    datasourceNew=driver.CreateDataSource(new_path)

    datasourceNew.CopyLayer(layerMode, layerName)
    # layerNew=datasourceNew.CreateLayer(layerName,spatialrefOld,geomType)
    # layerNew.CreateFields(schema)
    # layerOld.Erase(layerMode,layerNew)

    # feature.SetGeometry(geometry)
    # print(feature.GetGeometryRef().ExportToWkt())
    # layerNew.CreateFeature(feature)
    datasourceNew.SyncToDisk()
    datasourceNew.FlushCache()

    datasourceOld.Destroy()
    datasourceMode.Destroy()
    datasourceNew.Destroy()
    return

def vector_crstransform(*args):
    shapeFilePath,dstSrsESPG,destDataPath=args
    ogr.RegisterAll()
    gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
    data = ogr.Open(shapeFilePath)
    layer = data.GetLayer()
    spatial = layer.GetSpatialRef()
    layerName = layer.GetName()
    data.Destroy()
    dstSRS = osr.SpatialReference()
    dstSRS.ImportFromEPSG(int(dstSrsESPG))
    format = "ESRI Shapefile"
    options = gdal.VectorTranslateOptions(
		format=format,
		accessMode=None,
		srcSRS=spatial,
		dstSRS=dstSRS,
		reproject=True,
		selectFields=None,
		layerName=layerName,
		geometryType="POLYGON",
		dim="XY"
	)

    gdal.VectorTranslate(
        destDataPath,
        srcDS=shapeFilePath,
        options=options
    )
    return

def vector_2_raster(*args):
    shp_path, resolution, attribute, NoData_value, tif_path = args
    x_res, y_res = resolution
    datasource = ogr.Open(shp_path)
    layer = datasource.GetLayer()
    lon_min, lon_max, lat_min, lat_max = layer.GetExtent()
    s_projection = str(layer.GetSpatialRef())
    transform = (lon_min, x_res, 0, lat_max, 0, -y_res)

    width = int(abs((lon_max - lon_min) / x_res))
    height = int(abs((lat_max - lat_min) / y_res))

    driver = gdal.GetDriverByName('GTiff')
    target_ds = driver.Create(tif_path, width, height, 1, gdal.GDT_Byte)
    target_ds.SetGeoTransform(transform)
    target_ds.SetProjection(s_projection)
    options = [
        "ALL_TOUCHED=False",
        'MERGE_ALG=REPLACE',
        # "TILED=YES",
        # "COMPRESS=LZW",
        f"ATTRIBUTE={attribute}"]  # LZW,PACKBITS
    gdal.RasterizeLayer(target_ds, [1], layer,options=options)
    band = target_ds.GetRasterBand(1)
    band.SetNoDataValue(NoData_value)

    band.FlushCache()
    target_ds.FlushCache()

    target_ds = None
    del target_ds, layer
    return

def raster_pro_2_geo(*args):
    old_path, epsg, new_path = args
    datasetOLd = gdal.Open(old_path)
    width, height = datasetOLd.RasterXSize, datasetOLd.RasterYSize
    rasterCount = datasetOLd.RasterCount
    geotransform = datasetOLd.GetGeoTransform()
    projection = datasetOLd.GetProjection()

    spatialOld = osr.SpatialReference()
    spatialOld.ImportFromWkt(projection)

    spatialNew = osr.SpatialReference()
    spatialNew.ImportFromEPSG(epsg)

    transform  = osr.CoordinateTransformation(spatialOld, spatialNew)
    (ul_x, x_res, x_angle, ul_y, y_angle, y_res) = geotransform

    lr_x = ul_x + x_res*width + x_angle*height
    lr_y = ul_y + y_res*height + y_angle*width

    (new_ul_x, new_ul_y, _) = transform .TransformPoint(ul_y, ul_x)
    (new_lr_x, new_lr_y, _) = transform .TransformPoint(lr_y, lr_x)
   
    raster_type = gdal.GDT_Byte
    driver = gdal.GetDriverByName("GTiff")
    datasetNew = driver.Create(new_path,width,height,rasterCount,raster_type)
    # 转换后图像的分辨率
    x_size = (new_lr_x-new_ul_x)/width
    y_size = abs((new_lr_y-new_ul_y))/height
    # 转换后图像的六个放射变换参数
    geotransform = (new_ul_x, x_size, 0, new_ul_y, 0, -y_size)


    # datasetNew.SetGeoTransform(geotransform)
    datasetNew.SetProjection(spatialNew.ExportToWkt())
    # 投影转换后需要做重采样
    gdal.ReprojectImage(
        datasetOLd, 
        datasetNew, 
        src_wkt=spatialOld.ExportToWkt(), 
        dst_wkt=spatialNew.ExportToWkt(), 
        eResampleAlg=gdalconst.GRA_NearestNeighbour)

    # 关闭
    datasetOLd = None
    datasetNew= None
    return

def warp_pro2geo(*args):
    old_path, epsg, new_path = args
    datasetOLd = gdal.Open(old_path)
    projection = datasetOLd.GetProjection()

    spatialOld = osr.SpatialReference()
    spatialOld.ImportFromWkt(projection)

    spatialNew = osr.SpatialReference()
    spatialNew.ImportFromEPSG(epsg)

    options = gdal.WarpOptions(format='GTiff', srcSRS=spatialOld, dstSRS=spatialNew)
    gdal.Warp(new_path, old_path, options=options)
    return

def raster_compress(*args):
    old_path, method, new_path = args # LZW,PACKBITS
    dataset = gdal.Open(old_path)
    driver = gdal.GetDriverByName('GTiff')
    options = ["TILED=YES", f"COMPRESS={method}"]
    driver.CreateCopy(new_path, dataset, strict=1, options=options)
    del dataset
    return

def main():
    os.system(r"net use Z: \\192.168.2.172\root DELL123dell /user:root /PERSISTENT:YES")
    '''矢量坐标系转换''' 
    mode_path = r"Z:\projection\test\41河南省5级行政区22.01.21\41河南省xy.shp"
    newcrs_shp = r"Z:\projection\test\41河南省xy_.shp"
    if not os.path.exists(newcrs_shp):vector_crstransform(mode_path, 32650, newcrs_shp)
    '''矢量擦除'''
    old_path = r"Z:\projection\test\2022河南秋粮_玉米_总_原shp\2022河南秋粮_玉米_总.shp"
    fieldName = r"qhdm"
    fieldValue = r"410224"
    new_path = r"Z:\projection\test\2022河南秋粮_玉米_总_原shp\2022河南秋粮_玉米_总_.shp"
    vector_erase(old_path, newcrs_shp, fieldName, fieldValue, new_path)

    '''矢量合并'''
    
    '''矢量融合'''


    '''矢量转栅格'''
    shp_path = r"Z:\projection\test\2022河南秋粮_玉米_总_原shp\2022河南秋粮_玉米_总.shp"
    resolution = [3,3] 
    attribute = r"CLASS_ID" 
    NoData_value = 127
    tif_path = r"Z:\projection\test\test_compress.tif"
    # vector_2_raster(shp_path, resolution, attribute, NoData_value, tif_path)

    '''栅格坐标系转换'''
    old_path = r"Z:\projection\test\test_compress.tif" 
    epsg=4326
    new_path = r"Z:\projection\test\test_New.tif"
    # warp_pro2geo(old_path, epsg, new_path)

    '''栅格文件压缩'''
    old_path=r"Z:\projection\test\test_New.tif"
    method="LZW"
    new_path = r"Z:\projection\test\test_compressw.tif"
    # raster_compress(old_path, method, new_path)
    return

if __name__ == "__main__":
    main()

